package actors

import (
	"context"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/asynkron/protoactor-go/actor"
	"github.com/google/uuid"
	"github.com/sirupsen/logrus"

	"xagent/internal/memory"
	"xagent/internal/types"
	proto "xagent/proto/v1"
)

// BaseAgent provides common functionality for all agents
type BaseAgent struct {
	actor.Actor
	id             string
	name           string
	description    string
	state          string
	currentRole    types.Role
	actions        map[string]types.Action
	memory         memory.MemoryManager
	capabilities   []string
	context        actor.Context
	mutex          sync.RWMutex
	supervisor     *actor.PID
	logger         *logrus.Entry
	metrics        *AgentMetrics
	messageBatch   []*PrioritizedMessage
	batchStartTime time.Time
}

// AgentMetrics tracks agent performance metrics
type AgentMetrics struct {
	MessagesProcessed   int
	MessagesFailed      int
	AverageLatency      time.Duration
	LastActive          time.Time
	Status              string
	MessageQueueSize    int
	MemoryUsage         int64
	ProcessingTime      map[string]time.Duration
	ErrorCounts         map[string]int
	SuccessRate         float64
	ResourceUtilization float64
}

// Error types for supervision
type TemporaryError struct {
	Msg string
}

func (e *TemporaryError) Error() string {
	return e.Msg
}

type PermanentError struct {
	Msg string
}

func (e *PermanentError) Error() string {
	return e.Msg
}

// AgentState represents the possible states of an agent
type AgentState string

const (
	StateUninitialized AgentState = "uninitialized"
	StateIdle          AgentState = "idle"
	StateWorking       AgentState = "working"
	StatePaused        AgentState = "paused"
	StateError         AgentState = "error"
	StateRestarting    AgentState = "restarting"
	StateFinished      AgentState = "finished"
)

// StateTransition represents a valid state transition
type StateTransition struct {
	From      AgentState
	To        AgentState
	Condition func() bool
}

// validateStateTransition checks if a state transition is valid
func (a *BaseAgent) validateStateTransition(from, to AgentState) error {
	validTransitions := map[AgentState][]AgentState{
		StateUninitialized: {StateIdle},
		StateIdle:          {StateWorking, StatePaused, StateError},
		StateWorking:       {StateIdle, StateError, StatePaused},
		StatePaused:        {StateIdle, StateWorking, StateFinished},
		StateError:         {StateIdle, StateRestarting},
		StateRestarting:    {StateIdle, StateError},
		StateFinished:      {StateIdle},
	}

	if valid, exists := validTransitions[from]; exists {
		for _, validState := range valid {
			if validState == to {
				return nil
			}
		}
	}
	return fmt.Errorf("invalid state transition from %s to %s", from, to)
}

// NewBaseAgent creates a new base agent
func NewBaseAgent(name, description string) *BaseAgent {
	memMgr := memory.NewInMemoryManager()
	return &BaseAgent{
		id:           uuid.NewString(),
		name:         name,
		description:  description,
		state:        string(types.StateUninitialized),
		actions:      make(map[string]types.Action),
		capabilities: make([]string, 0),
		memory:       memMgr,
		logger:       logrus.WithFields(logrus.Fields{"agent": name}),
		metrics:      &AgentMetrics{LastActive: time.Now()},
	}
}

// Ensure BaseAgent implements types.Agent
var _ types.Agent = (*BaseAgent)(nil)

// GetID returns the agent's ID
func (a *BaseAgent) GetID() string {
	return a.id
}

// GetName returns the agent's name
func (a *BaseAgent) GetName() string {
	return a.name
}

// GetDescription returns the agent's description
func (a *BaseAgent) GetDescription() string {
	return a.description
}

// GetRole returns the agent's current role name
func (a *BaseAgent) GetRole() string {
	a.mutex.RLock()
	defer a.mutex.RUnlock()

	if a.currentRole == nil {
		return ""
	}
	return a.currentRole.GetName()
}

// AssignRole assigns a new role to the agent
func (a *BaseAgent) AssignRole(role types.Role) error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	if role == nil {
		return fmt.Errorf("cannot assign nil role")
	}

	// Revoke current role if exists
	if a.currentRole != nil {
		if err := a.currentRole.OnRevoke(a); err != nil {
			return fmt.Errorf("failed to revoke current role: %w", err)
		}
	}

	// Assign new role
	if err := role.OnAssign(a); err != nil {
		return fmt.Errorf("failed to assign new role: %w", err)
	}

	a.currentRole = role
	a.logger = a.logger.WithField("role", role.GetName())
	return nil
}

// GetCurrentRole returns the agent's current role
func (a *BaseAgent) GetCurrentRole() types.Role {
	a.mutex.RLock()
	defer a.mutex.RUnlock()
	return a.currentRole
}

// AddAction adds an action to the agent
func (a *BaseAgent) AddAction(action types.Action) error {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	if action == nil {
		return fmt.Errorf("cannot add nil action")
	}

	if _, exists := a.actions[action.GetName()]; exists {
		return fmt.Errorf("action %s already exists", action.GetName())
	}

	a.actions[action.GetName()] = action
	return nil
}

// GetActions returns all actions available to the agent
func (a *BaseAgent) GetActions() []types.Action {
	a.mutex.RLock()
	defer a.mutex.RUnlock()

	actions := make([]types.Action, 0, len(a.actions))
	for _, action := range a.actions {
		actions = append(actions, action)
	}
	return actions
}

// ExecuteAction executes a named action with the given input
func (a *BaseAgent) ExecuteAction(ctx context.Context, actionName string, input interface{}) (*proto.AgentMessage, error) {
	a.mutex.RLock()
	action, exists := a.actions[actionName]
	a.mutex.RUnlock()

	if !exists {
		return nil, fmt.Errorf("action %s not found", actionName)
	}

	return action.Execute(ctx, input)
}

// ResourceStats tracks resource usage
type ResourceStats struct {
	CPUUsage    float64
	MemoryUsage int64
	TaskQueue   int
	LoadFactor  float64
}

// ResourceConfig defines resource limits
type ResourceConfig struct {
	MaxCPUUsage    float64
	MaxMemoryUsage int64
	MaxTaskQueue   int
	MaxLoadFactor  float64
}

// updateResourceStats updates resource usage statistics
func (a *BaseAgent) updateResourceStats() {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	stats := &ResourceStats{
		CPUUsage:    a.metrics.ResourceUtilization,
		MemoryUsage: a.metrics.MemoryUsage,
		TaskQueue:   a.metrics.MessageQueueSize,
		LoadFactor:  float64(a.metrics.MessageQueueSize) / float64(defaultMaxTaskQueue),
	}

	// Update metrics
	a.metrics.ResourceUtilization = stats.CPUUsage
	a.metrics.MemoryUsage = stats.MemoryUsage
	a.metrics.MessageQueueSize = stats.TaskQueue

	// Log if approaching resource limits
	if stats.LoadFactor > 0.8 {
		a.logger.WithFields(logrus.Fields{
			"load_factor": stats.LoadFactor,
			"task_queue":  stats.TaskQueue,
		}).Warn("High load detected")
	}
}

// canAcceptTask checks if the agent can accept a new task
func (a *BaseAgent) canAcceptTask() bool {
	a.mutex.RLock()
	defer a.mutex.RUnlock()

	return a.metrics.MessageQueueSize < defaultMaxTaskQueue &&
		a.metrics.ResourceUtilization < 0.9 &&
		types.State(a.state) == types.StateIdle
}

// ProcessTask with resource management
func (a *BaseAgent) ProcessTask(task *proto.Task) error {
	if !a.canAcceptTask() {
		return &TemporaryError{Msg: "Agent is currently overloaded"}
	}

	a.SetState(types.StateWorking)
	defer a.SetState(types.StateIdle)

	a.logger.WithFields(logrus.Fields{
		"task_id": task.Id,
		"type":    task.Name,
	}).Info("Processing task")

	// Execute task as action if exists
	if action, exists := a.actions[task.Name]; exists {
		result, err := action.Execute(context.Background(), task)
		if err != nil {
			a.handleError(err)
			return err
		}
		if a.context != nil && a.context.Parent() != nil {
			a.context.Send(a.context.Parent(), result)
		}
	}

	return nil
}

// MessagePriority defines message priority levels
type MessagePriority int

const (
	LowPriority MessagePriority = iota
	MediumPriority
	HighPriority
)

// PrioritizedMessage wraps a message with priority information
type PrioritizedMessage struct {
	Message  *proto.AgentMessage
	Priority MessagePriority
	Created  time.Time
}

// determineMessagePriority assigns priority based on message type
func (a *BaseAgent) determineMessagePriority(msg *proto.AgentMessage) MessagePriority {
	switch msg.Type {
	case proto.MessageType_MESSAGE_TYPE_ERROR:
		return HighPriority
	case proto.MessageType_MESSAGE_TYPE_STATUS:
		return MediumPriority
	default:
		return LowPriority
	}
}

// HandleMessage with priority handling
func (a *BaseAgent) HandleMessage(msg *proto.AgentMessage) error {
	priority := a.determineMessagePriority(msg)
	pMsg := &PrioritizedMessage{
		Message:  msg,
		Priority: priority,
		Created:  time.Now(),
	}

	a.mutex.Lock()
	if len(a.messageBatch) == 0 {
		a.batchStartTime = time.Now()
	}
	a.messageBatch = append(a.messageBatch, pMsg)
	shouldProcess := len(a.messageBatch) >= 10 ||
		priority == HighPriority ||
		time.Since(a.batchStartTime) > 5*time.Second
	a.mutex.Unlock()

	if shouldProcess {
		return a.processBatch()
	}

	return nil
}

// processBatch processes all messages in the current batch
func (a *BaseAgent) processBatch() error {
	a.mutex.Lock()
	batch := a.messageBatch
	a.messageBatch = make([]*PrioritizedMessage, 0)
	a.mutex.Unlock()

	// Sort by priority
	sort.Slice(batch, func(i, j int) bool {
		return batch[i].Priority > batch[j].Priority
	})

	for _, pMsg := range batch {
		if err := a.processMessage(pMsg.Message); err != nil {
			a.handleError(err)
			continue
		}
	}

	return nil
}

// processMessage handles a single message
func (a *BaseAgent) processMessage(msg *proto.AgentMessage) error {
	entry := &memory.MemoryEntry{
		ID:      uuid.NewString(),
		Content: a.getMessageContent(msg),
		Type:    memory.MemoryTypeConversation,
		Metadata: map[string]interface{}{
			"sender":   msg.Sender,
			"type":     msg.Type.String(),
			"priority": a.determineMessagePriority(msg),
		},
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}

	return a.memory.Store(context.Background(), entry)
}

// GetState returns the agent's current state
func (a *BaseAgent) GetState() types.State {
	a.mutex.RLock()
	defer a.mutex.RUnlock()
	return types.State(a.state)
}

// SetState sets the agent's state
func (a *BaseAgent) SetState(newState types.State) {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	// Validate state transition
	if err := a.validateStateTransition(AgentState(a.state), AgentState(newState)); err != nil {
		a.logger.WithError(err).Error("Invalid state transition")
		return
	}

	prevState := a.state
	a.state = string(newState)

	// Notify supervisor of state change
	if a.supervisor != nil {
		stateMsg := &proto.AgentMessage{
			Id:        uuid.NewString(),
			Type:      proto.MessageType_MESSAGE_TYPE_STATUS,
			Sender:    a.id,
			Timestamp: time.Now().Unix(),
			Content: &proto.AgentMessage_Message{
				Message: fmt.Sprintf("Agent %s state changed from %s to %s",
					a.id, prevState, newState),
			},
		}
		a.context.Send(a.supervisor, stateMsg)
	}

	a.logger.WithFields(logrus.Fields{
		"previous_state": prevState,
		"new_state":      newState,
	}).Info("Agent state changed")
}

// Helper methods

func (a *BaseAgent) handleError(err error) {
	a.metrics.MessagesFailed++
	a.logger.WithError(err).Error("Error occurred")

	if a.context != nil && a.context.Parent() != nil {
		errMsg := &proto.AgentMessage{
			Id:   uuid.NewString(),
			Type: proto.MessageType_MESSAGE_TYPE_ERROR,
			Content: &proto.AgentMessage_Error{
				Error: &proto.Error{
					Code:    "INTERNAL_ERROR",
					Message: err.Error(),
				},
			},
		}
		a.context.Send(a.context.Parent(), errMsg)
	}
}

func (a *BaseAgent) getMessageContent(msg *proto.AgentMessage) string {
	switch c := msg.Content.(type) {
	case *proto.AgentMessage_Message:
		return c.Message
	case *proto.AgentMessage_Error:
		return c.Error.Message
	default:
		return ""
	}
}

// Receive handles incoming messages
func (a *BaseAgent) Receive(context actor.Context) {
	startTime := time.Now()
	a.context = context

	switch context.Message().(type) {
	case *actor.Started:
		a.OnStarted()
	case *actor.Stopping:
		a.OnStopping()
	case *actor.Stopped:
		a.OnStopped()
	case *actor.Restarting:
		a.OnRestarting(context.Message())
	case *proto.AgentMessage:
		if err := a.HandleMessage(context.Message().(*proto.AgentMessage)); err != nil {
			a.handleError(err)
		}
	case *proto.Task:
		if err := a.ProcessTask(context.Message().(*proto.Task)); err != nil {
			a.handleError(err)
		}
	}

	a.updateMetrics(startTime)
}

// Lifecycle hooks
func (a *BaseAgent) OnStarted() {
	a.SetState(types.StateIdle)
	a.logger.Info("Agent started")
}

func (a *BaseAgent) OnStopping() {
	a.SetState(types.StatePaused)
	a.logger.Info("Agent stopping")
}

func (a *BaseAgent) OnStopped() {
	a.SetState(types.StateFinished)
	a.logger.Info("Agent stopped")
}

func (a *BaseAgent) OnRestarting(reason interface{}) {
	a.handleRestart(reason)
}

func (a *BaseAgent) updateMetrics(startTime time.Time) {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	processingTime := time.Since(startTime)
	a.metrics.LastActive = time.Now()
	a.metrics.MessagesProcessed++

	// Update processing time for different message types
	if a.metrics.ProcessingTime == nil {
		a.metrics.ProcessingTime = make(map[string]time.Duration)
	}

	// Calculate success rate
	total := float64(a.metrics.MessagesProcessed + a.metrics.MessagesFailed)
	if total > 0 {
		a.metrics.SuccessRate = float64(a.metrics.MessagesProcessed) / total * 100
	}

	// Update average latency
	if a.metrics.AverageLatency == 0 {
		a.metrics.AverageLatency = processingTime
	} else {
		a.metrics.AverageLatency = (a.metrics.AverageLatency + processingTime) / 2
	}
}

// GetCapabilities returns agent's capabilities
func (a *BaseAgent) GetCapabilities() []string {
	a.mutex.RLock()
	defer a.mutex.RUnlock()
	return append([]string{}, a.capabilities...)
}

// AddCapability adds a capability to the agent
func (a *BaseAgent) AddCapability(capability string) {
	a.mutex.Lock()
	defer a.mutex.Unlock()
	a.capabilities = append(a.capabilities, capability)
	a.logger.WithField("capability", capability).Info("Added new capability")
}

// GetContext returns the actor context
func (a *BaseAgent) GetContext() actor.Context {
	return a.context
}

// GetSupervisor returns the actor's supervisor
func (a *BaseAgent) GetSupervisor() *actor.PID {
	return a.supervisor
}

// SetSupervisor sets the actor's supervisor
func (a *BaseAgent) SetSupervisor(supervisor *actor.PID) {
	a.supervisor = supervisor
	a.logger.WithField("supervisor", supervisor).Info("Supervisor set")
}

// GetMetrics returns agent metrics
func (a *BaseAgent) GetMetrics() *AgentMetrics {
	return a.metrics
}

// SendMessage sends a message to another agent
func (a *BaseAgent) SendMessage(target *actor.PID, msg *proto.AgentMessage) {
	a.context.Send(target, msg)
	a.logger.WithFields(logrus.Fields{
		"target": target,
		"type":   msg.Type,
	}).Debug("Message sent")
}

// SupervisionStrategy returns the supervision strategy for the agent
func (a *BaseAgent) SupervisionStrategy() actor.SupervisorStrategy {
	return actor.NewOneForOneStrategy(10, 10*time.Second, func(reason interface{}) actor.Directive {
		a.logger.Errorf("Actor failed with reason: %v", reason)

		switch reason.(type) {
		case *TemporaryError:
			return actor.RestartDirective
		case *PermanentError:
			return actor.StopDirective
		default:
			if a.metrics.MessagesFailed < 5 {
				return actor.ResumeDirective
			}
			return actor.RestartDirective
		}
	})
}

// handleRestart handles cleanup when actor is restarting
func (a *BaseAgent) handleRestart(reason interface{}) {
	a.logger.Warnf("Agent %s is restarting due to: %v", a.id, reason)
	a.metrics.MessagesFailed++
	a.metrics.Status = "restarting"
	a.SetState(types.StateRestarting)

	// Notify supervisor
	if a.supervisor != nil {
		stateMsg := &proto.AgentMessage{
			Id:        uuid.NewString(),
			Type:      proto.MessageType_MESSAGE_TYPE_STATUS,
			Sender:    a.id,
			Timestamp: time.Now().Unix(),
		}
		stateMsg.Content = &proto.AgentMessage_Message{
			Message: fmt.Sprintf("Agent %s state changed to restarting: %v", a.id, reason),
		}
		a.context.Send(a.supervisor, stateMsg)
	}

	// Clean up resources
	a.cleanupResources()
}

// cleanupResources handles resource cleanup during restart/shutdown
func (a *BaseAgent) cleanupResources() {
	a.mutex.Lock()
	defer a.mutex.Unlock()

	// Clear any pending tasks
	a.actions = make(map[string]types.Action)

	// Reset metrics
	a.metrics.MessageQueueSize = 0
	a.metrics.MemoryUsage = 0

	// Clear memory if needed
	if a.memory != nil {
		a.memory.Clear(context.Background())
	}
}

const (
	defaultMaxTaskQueue = 100
)
